"""
重构后的报表分析工作流 - 使用新的管理器架构
"""
import json
from typing import Dict, Any, Optional, AsyncGenerator

from autogen_agentchat.messages import TextMessage
from autogen_core import CancellationToken

from core.workflow_manager import workflow_manager
from core.agent_manager import agent_manager
from core.mcp_manager import mcp_manager
from core.prompt_manager import prompt_manager
from config.yaml_settings import get_llm_model_config, get_llm_api_config, LLMModelConfig, LLMAPIConfig


class ReportAnalysisWorkflowV2:
    """重构后的报表分析工作流协调器"""

    def __init__(self,
                 api_key: Optional[str] = None,
                 base_url: Optional[str] = None,
                 model_config: Optional[LLMModelConfig] = None,
                 api_config: Optional[LLMAPIConfig] = None):
        # 使用传入的配置或默认配置
        if api_config is None:
            api_config = get_llm_api_config()
        if model_config is None:
            model_config = get_llm_model_config()

        self.api_key = api_key or api_config.api_key
        self.base_url = base_url or api_config.base_url
        self.model_config = model_config
        self.api_config = api_config
        self._initialized = False

    async def initialize(self):
        """初始化工作流组件"""
        if self._initialized:
            return

        try:
            # 初始化MCP工具（单例）
            await mcp_manager.get_sql_tools()
            
            # 初始化Agent（单例）
            agent_manager.get_analysis_agent(
                api_key=self.api_key,
                base_url=self.base_url,
                model_config=self.model_config,
                api_config=self.api_config
            )
            
            agent_manager.get_chart_agent(
                api_key=self.api_key,
                base_url=self.base_url,
                model_config=self.model_config,
                api_config=self.api_config
            )

            self._initialized = True
            print("重构后的报表分析工作流初始化完成")

        except Exception as e:
            print(f"工作流初始化失败: {str(e)}")
            raise

    async def analyze_report_data(self, user_query: str) -> Dict[str, Any]:
        """执行报表数据分析工作流 - 非流式版本"""
        if not self._initialized:
            await self.initialize()

        cancellation_token = CancellationToken()

        try:
            print(f"开始处理用户查询: {user_query}")

            # 步骤1: 获取数据库表结构信息
            print("步骤1: 获取数据库表结构...")
            sql_tools = await mcp_manager.get_sql_tools()
            table_info_result = await sql_tools.get_all_tables(cancellation_token)
            
            if table_info_result.get("status") == "error":
                return {
                    "status": "error",
                    "user_query": user_query,
                    "error": table_info_result.get("error", "获取表结构失败")
                }
            
            table_info = table_info_result.get("data", "")
            print(f"表结构信息: {table_info}")

            # 步骤2: LLM分析用户意图并生成SQL
            print("步骤2: LLM分析用户意图并生成SQL...")
            sql_query = await self._generate_sql_query(user_query, table_info, cancellation_token)
            if sql_query.startswith("SQL生成失败"):
                return {
                    "status": "error",
                    "user_query": user_query,
                    "error": sql_query
                }
            print(f"生成的SQL: {sql_query}")

            # 步骤3: 执行SQL查询
            print("步骤3: 执行SQL查询...")
            query_result = await sql_tools.execute_custom_sql(sql_query, cancellation_token)
            if query_result.get("status") == "error":
                return {
                    "status": "error",
                    "user_query": user_query,
                    "error": query_result.get("error", "SQL执行失败")
                }
            print(f"查询结果: {query_result}")

            # 步骤4: 生成HTML报告
            print("步骤4: 生成HTML报告...")
            html_report = await self._generate_html_report(query_result, user_query, cancellation_token)
            if html_report.startswith("HTML生成失败"):
                return {
                    "status": "error",
                    "user_query": user_query,
                    "error": html_report
                }

            # 步骤5: 返回结果
            result = {
                "status": "success",
                "user_query": user_query,
                "sql_query": sql_query,
                "data": query_result.get("data", ""),
                "html_report": html_report
            }

            print("报表分析工作流执行完成")
            return result

        except Exception as e:
            error_result = {
                "status": "error",
                "user_query": user_query,
                "error": str(e)
            }
            print(f"工作流执行失败: {str(e)}")
            return error_result

    async def _generate_sql_query(self, user_query: str, table_info: str, cancellation_token: CancellationToken) -> str:
        """LLM生成SQL查询 - 使用提示词模板"""
        try:
            # 从提示词管理器加载模板
            prompt_template = prompt_manager.load_prompt("sql_generation")
            
            # 格式化提示词
            prompt = prompt_template.format(
                user_query=user_query,
                table_info=table_info
            )
            
            # 使用Agent管理器获取分析Agent
            analysis_agent = agent_manager.get_analysis_agent(
                api_key=self.api_key,
                base_url=self.base_url,
                model_config=self.model_config,
                api_config=self.api_config
            )
            
            response = await analysis_agent.on_messages(
                [TextMessage(content=prompt, source="workflow")],
                cancellation_token
            )
            return response.chat_message.content.strip()
        except Exception as e:
            return f"SQL生成失败: {str(e)}"

    async def _generate_html_report(self, query_result: Dict[str, Any], user_query: str,
                                    cancellation_token: CancellationToken) -> str:
        """生成完整的HTML报告 - 使用提示词模板"""
        try:
            # 提取查询结果数据
            data = query_result.get("data", [])
            if not data:
                return "HTML生成失败: 查询结果为空"
            
            # 数据处理逻辑（与原版本相同）
            if isinstance(data, str):
                try:
                    data = json.loads(data)
                except:
                    try:
                        data = eval(data)
                    except:
                        pass
            
            if not isinstance(data, list):
                data = [data]
            
            # 分析数据结构
            columns = []
            data_types = {}
            
            if data and isinstance(data[0], dict):
                columns = list(data[0].keys())
                
                for col in columns:
                    data_types[col] = "string"
                    sample_size = min(10, len(data))
                    numeric_count = 0
                    date_count = 0
                    
                    for i in range(sample_size):
                        if i < len(data) and col in data[i]:
                            value = data[i][col]
                            
                            if isinstance(value, (int, float)) or (
                                    isinstance(value, str) and value.replace('.', '', 1).isdigit()):
                                numeric_count += 1
                            
                            if isinstance(value, str) and (
                                    '-' in value or '/' in value) and (
                                    len(value) >= 8 and len(value) <= 30):
                                date_count += 1
                    
                    if numeric_count >= sample_size * 0.7:
                        data_types[col] = "numeric"
                    elif date_count >= sample_size * 0.7:
                        data_types[col] = "date"
            
            # 准备数据样例
            try:
                sample_data_json = json.dumps(data[:3], ensure_ascii=False, indent=2)
            except TypeError:
                sample_data = []
                for i, item in enumerate(data[:3]):
                    if i >= 3:
                        break
                    try:
                        if isinstance(item, dict):
                            sample_item = {}
                            for k, v in item.items():
                                if isinstance(v, (str, int, float, bool, type(None))):
                                    sample_item[k] = v
                                else:
                                    sample_item[k] = str(v)
                            sample_data.append(sample_item)
                        else:
                            sample_data.append(str(item))
                    except:
                        sample_data.append(f"[复杂数据项 {i+1}]")
                
                sample_data_json = str(sample_data)
            
            # 从提示词管理器加载模板
            prompt_template = prompt_manager.load_prompt("chart_generation")
            
            # 格式化提示词
            prompt = prompt_template.format(
                user_query=user_query,
                data_count=len(data),
                columns=columns,
                data_types=data_types,
                sample_data=sample_data_json
            )
            
            # 使用Agent管理器获取图表Agent
            chart_agent = agent_manager.get_chart_agent(
                api_key=self.api_key,
                base_url=self.base_url,
                model_config=self.model_config,
                api_config=self.api_config
            )
            
            print(f"[HTML生成] 发送请求到LLM...")
            
            response = await chart_agent.on_messages(
                [TextMessage(content=prompt, source="workflow")],
                cancellation_token
            )
            
            html_report = response.chat_message.content.strip()
            print(f"[HTML生成] HTML报告生成成功，长度: {len(html_report)}")
            
            # 清理可能的Markdown代码块标记
            if html_report.startswith("```html"):
                html_report = html_report.replace("```html", "", 1)
                if html_report.endswith("```"):
                    html_report = html_report[:-3]
            
            html_report = html_report.strip()
            return html_report
            
        except Exception as e:
            error_msg = f"HTML生成失败: {str(e)}"
            print(f"[HTML生成] 错误: {error_msg}")
            return error_msg

    async def get_workflow_status(self) -> Dict[str, Any]:
        """获取工作流状态"""
        if not self._initialized:
            await self.initialize()
        
        return {
            "status": "ready" if self._initialized else "not_initialized",
            "available_tools": ["get_all_tables", "get_table_structure", "execute_custom_sql"],
            "message": "重构后的工作流已初始化" if self._initialized else "工作流未初始化",
            "agent_info": agent_manager.get_agent_info(),
            "mcp_initialized": mcp_manager.is_initialized()
        }